/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package network;

/**
 *
 * @author Raise
 */
import network.Layer;
import java.util.Random;

public class ConvLayer extends Layer {

    private double[][][] kernels;  //卷积核维度：[滤波器数, 卷积核高, 卷积核宽],单输入
    private double[] biases;//偏置
    private double[][][] lastInput;//上次输入 [高, 宽, 通道] - 用于反向传播
    private int kernelSize;// 卷积核尺寸 - 通常为3,5,7等奇数
    private int numFilters;// 滤波器数量 - 决定输出通道数（其实就是卷积核数量）
    private int inputHeight;// 输入尺寸
    private int inputWidth;
    private int outputHeight;  // 输出高度
    private int outputWidth;   // 输出宽度
    private int outputSize;    // 输出总元素数
    private double[] lastSums;
    private double[][][] kernelGradients;
    private double[] biasGradients;

    public ConvLayer(int inputHeight, int inputWidth, int kernelSize, int numFilters, Random rand) {
        this.inputHeight = inputHeight;
        this.inputWidth = inputWidth;
        this.kernelSize = kernelSize;
        this.numFilters = numFilters;
        this.inputSize = inputHeight * inputWidth;

        initializeParameters(rand);
    }

    private void initializeParameters(Random rand) {
        //Random rand = new Random();
        kernels = new double[numFilters][kernelSize][kernelSize];
        biases = new double[numFilters];

        // Xavier初始化
        double scale = Math.sqrt(2.0 / (kernelSize * kernelSize));
        for (int f = 0; f < numFilters; f++) {
            for (int i = 0; i < kernelSize; i++) {
                for (int j = 0; j < kernelSize; j++) {
                    kernels[f][i][j] = rand.nextGaussian() * scale;
                }
            }
            biases[f] = 0.1;
        }
    }

    @Override
    public double[] forward(double[] input) {
        // 将一维输入转换为三维 [1][height][width] (单通道)
        lastInput = reshapeInput(input);

        // 输出尺寸计算
        int outputHeight = inputHeight - kernelSize + 1;
        int outputWidth = inputWidth - kernelSize + 1;
        this.outputSize = outputHeight * outputWidth * numFilters;
        double[] output = new double[outputSize];//输出还是一维的
        this.lastSums = new double[outputSize];  // 保存加权和用于反向传播

        int outputIndex = 0;
        for (int f = 0; f < numFilters; f++) {
            for (int i = 0; i < outputHeight; i++) {
                for (int j = 0; j < outputWidth; j++) {
                    double sum = biases[f];
                    // 卷积运算
                    for (int ki = 0; ki < kernelSize; ki++) {
                        for (int kj = 0; kj < kernelSize; kj++) {
                            sum += lastInput[0][i + ki][j + kj] * kernels[f][ki][kj];
                        }
                    }
                    lastSums[outputIndex] = sum;           // 保存加权和
                    output[outputIndex] = Math.max(0, sum); // ReLU激活
                    outputIndex++;
                }
            }
        }

        return output;
    }

    @Override
    public double[] backward(double[] upstreamGradient, double learningRate) {

        
        // 1.1 将一维的上游梯度转换为三维结构
        int outputHeight = inputHeight - kernelSize + 1;
        int outputWidth = inputWidth - kernelSize + 1;
        double[][][] upstream3D = reshapeGradient3D(upstreamGradient, outputHeight, outputWidth, numFilters);

        // 1.2 初始化梯度数组
//        double[][][] kernelGradients = new double[numFilters][kernelSize][kernelSize];
//        double[] biasGradients = new double[numFilters];
        this.kernelGradients = new double[numFilters][kernelSize][kernelSize];  // ✅ 改为this
        this.biasGradients = new double[numFilters];  // ✅ 改为this
        double[][][] inputGradient3D = new double[1][inputHeight][inputWidth]; // 初始化为0

        // 2.1 将lastSums（前向传播保存的加权和）转换为三维
        double[][][] sums3D = reshapeGradient3D(lastSums, outputHeight, outputWidth, numFilters);

        // 2.2 计算考虑ReLU导数的delta
        double[][][] delta = new double[numFilters][outputHeight][outputWidth];
        for (int f = 0; f < numFilters; f++) {
            for (int i = 0; i < outputHeight; i++) {
                for (int j = 0; j < outputWidth; j++) {
                    // δ = upstream_gradient × ReLU_derivative
                    delta[f][i][j] = upstream3D[f][i][j] * (sums3D[f][i][j] > 0 ? 1.0 : 0.0);
                }
            }
        }

        // 3.计算权重梯度，损失函数对每个卷积核权重的偏导数
        for (int f = 0; f < numFilters; f++) {
            for (int ki = 0; ki < kernelSize; ki++) {
                for (int kj = 0; kj < kernelSize; kj++) {
                    double grad = 0;

                    // 累加所有使用到这个权重的位置
                    for (int i = 0; i < outputHeight; i++) {
                        for (int j = 0; j < outputWidth; j++) {
                            grad += delta[f][i][j] * lastInput[0][i + ki][j + kj];
                        }
                    }
                    kernelGradients[f][ki][kj] = grad;
                }
            }
        }

        //4.计算偏置梯度
        for (int f = 0; f < numFilters; f++) {
            double biasGrad = 0;
            for (int i = 0; i < outputHeight; i++) {
                for (int j = 0; j < outputWidth; j++) {
                    biasGrad += delta[f][i][j];
                }
            }
            biasGradients[f] = biasGrad;
        }

        //5：计算输入梯度（传递给前一层）
        for (int f = 0; f < numFilters; f++) {
            for (int i = 0; i < outputHeight; i++) {
                for (int j = 0; j < outputWidth; j++) {
                    for (int ki = 0; ki < kernelSize; ki++) {
                        for (int kj = 0; kj < kernelSize; kj++) {
                            int inputY = i + ki;
                            int inputX = j + kj;
                            inputGradient3D[0][inputY][inputX] += delta[f][i][j] * kernels[f][ki][kj];
                        }
                    }
                }
            }
        }

        // 6.更新卷积核权重
        for (int f = 0; f < numFilters; f++) {
            for (int ki = 0; ki < kernelSize; ki++) {
                for (int kj = 0; kj < kernelSize; kj++) {
                    kernels[f][ki][kj] -= learningRate * kernelGradients[f][ki][kj];
                }
            }
        }

        // 更新偏置
        for (int f = 0; f < numFilters; f++) {
            biases[f] -= learningRate * biasGradients[f];
        }

        return flatten3D(inputGradient3D);
    }

//    将一维输入数组重塑为三维张量，适配卷积层输入格式(前向传播要用)
//    在MNIST数据集上使用：将784个像素的一维数组转换为28×28的二维图像格式
    private double[][][] reshapeInput(double[] input) {
        // 创建三维张量：批次大小=1，因为每次处理单张图片
        double[][][] reshaped = new double[1][inputHeight][inputWidth];

        // 遍历所有像素位置，将一维数组映射到二维网格
        for (int i = 0; i < inputHeight; i++) {
            for (int j = 0; j < inputWidth; j++) {
                // 计算一维数组中的索引：行索引×宽度 + 列索引
                // 将展平的像素数据恢复为图像的空间结构
                reshaped[0][i][j] = input[i * inputWidth + j];
            }
        }
        return reshaped;
    }

    /**
     * 将一维梯度数组转换为三维结构 [numFilters][height][width]
     */
    private double[][][] reshapeGradient3D(double[] flat, int height, int width, int numFilters) {
        double[][][] result = new double[numFilters][height][width];
        int index = 0;
        for (int f = 0; f < numFilters; f++) {
            for (int i = 0; i < height; i++) {
                for (int j = 0; j < width; j++) {
                    result[f][i][j] = flat[index++];
                }
            }
        }
        return result;
    }

    /**
     * 将三维数组展平为一维
     */
    private double[] flatten3D(double[][][] array3D) {
        int channels = array3D.length;
        int height = array3D[0].length;
        int width = array3D[0][0].length;
        double[] flat = new double[channels * height * width];

        int index = 0;
        for (int c = 0; c < channels; c++) {
            for (int i = 0; i < height; i++) {
                for (int j = 0; j < width; j++) {
                    flat[index++] = array3D[c][i][j];
                }
            }
        }
        return flat;
    }

    /**
     * 获取权重的梯度（用于测试）
     */
    public double getWeightGradient(int f, int ki, int kj) {
        if (this.kernelGradients == null) {
            throw new IllegalStateException("请先调用backward方法计算梯度");
        }
        return this.kernelGradients[f][ki][kj];
    }

    /**
     * 获取权重（用于测试）
     */
    public double getWeight(int f, int ki, int kj) {
        return kernels[f][ki][kj];
    }

    /**
     * 设置权重（用于测试）
     */
    public void setWeight(int f, int ki, int kj, double value) {
        kernels[f][ki][kj] = value;
    }
}
